#!/usr/bin/env python

import socket
import time
import sys
import threading
import re

class Kismet(threading.Thread):
    
    def __init__(self, host = "localhost", port = 2501):
        threading.Thread.__init__(self)
        self.setDaemon(True)
        self.port = port
        self.host = host
        self.die = 0
        self.cmd = 0
        self.parseProtocolTree = True
        self.protocols = [ ]
        self.callbacks = { }
        self.ackbacks = { }
        self.dangle = None # Used for line splitting on the socket buffer
        
        self.connect()
        
    def __str__(self):
        return """
Listening to Kismet version %s, on server %s.
Server started on %s. UID: %i.
Dump files: %s.
""" % (self.version,
       self.servername,
       time.strftime('%c' ,self.starttime),
       self.uid,
       self.dumpfiles)

    def connect(self):
        self.die = 0
        
        self.conn = None
        try:
            self.conn = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.conn.connect((self.host, self.port))
            
            # Get Kismet Version Info and Create Protocol Tree:
            self.getKismetMessage()
            
        except socket.error, msg:
            self.conn = None
            print msg
            
        if self.conn is None:
            print 'Could not connect to kismet'
            sys.exit(1)

    def run(self):
        while True:
            # print self.conn.recv(4096) + "\n"
            if self.die: break
            try:
                self.getKismetMessage()
            except Exception, e:
                print e
                break
            
    def getKismetMessage(self):
        try:
            # Apparently ruby's conn.get only gives back one line at a time... not so here.
            # self.parseline(self.conn.recv(4096))
            data = self.conn.recv(4096)
            data = data.split('\n')
            if self.dangle: # If there was a piece of a line left over from the last buffer
            	data[0] = self.dangle+data[0] # Prepend it to the first line of this buffer
            	self.dangle = None # Then remove the piece
            if data[-1]!='': # If this buffer doesn't end with a full line, the last item in data will *not* be blank
            	self.dangle = data[-1] # In that case, save it for prepending to the next buffer
            	del data[-1] # Remove that piece from data, so we don't parseline() an incomplete line
            else:
            	self.dangle = None # If we happen to get a full line at the end of the buffer.
            for line in data:
                self.parseline(line)
        except socket.error, msg: 
            print "Receive error: %s" % msg

    def sendraw(self, txd):
        try:
            self.conn.send(txd)
        except socket.error, msg:
            print "Send error: %s" % msg
    
    def kill(self):
        self.die = 1
        self.conn.close()

    def wait(self):
        self.join()

    def parseline(self, line):
        pattern = re.compile("\*([A-Z0-9]+): (.*)")

        md = pattern.match(line)
        
        # DEBUG TIP: Original ruby required exit? exit if md == nil
        if md == None: return
        
        # print md.length

        # for ss in md.groups():
        #    print ss

        if len(md.groups()) != 2: return

        p = md.group(1).upper()
        
        if p == "ERROR":
            f = self.parsedata(["cmdid","text"], md.group(2))
            raise KismetError(f['cmdid'], f['text'])
            
        elif p == "TIME":
            f = self.parsedata(["timesec"], md.group(2))
            self.timestamp = time.localtime(int(f["timesec"]))
            
        elif p == "ACK":
            
            f = self.parsedata(["cmdid", "text"], md.group(2))
            id = int(f['cmdid'])

            if id in self.ackbacks:
                self.ackbacks[id](id, f["text"])
                self.ackbacks.pop(id)

            return p

        elif p in self.callbacks:
            f = self.parsedata(self.callbacks[p][0], md.group(2))
            # print "%s got %i fields" % (p, len(f.keys()))
            self.callbacks[p][1](p, f)
            
        elif p == "KISMET":
            f = self.parsedata(['version', 'starttime', 'servername', 'dumpfiles', 'uid'], md.group(2))             
            self.version = f['version']
            self.starttime = time.localtime(int(f['starttime']))
            self.servername = f['servername']
            self.dumpfiles = f['dumpfiles']
            self.uid = int(f['uid'])
            
        elif p == "PROTOCOLS":
            f = self.parsedata(['protocols'], md.group(2))
            if self.parseProtocolTree:
                protocols = f['protocols'].split(',')
                self.getKismetProtocolTree(protocols)
                
        elif p == "CAPABILITY":
            f = self.parsedata(['protocol', 'capabilities'], md.group(2))
            self.protocols.append({'protocol':f['protocol'], 'capabilities':f['capabilities'].split(',')})

        return md.group(2)

    def parsedata(self, fields, data):        
        # Append a space at end of data to signal algorithm the end has been reached:
        data = data + ' '
        
        in_delim = 0

        da = {}
        f = ""
        fnum = 0

        for char in data:
            if char == "\001" and in_delim == 1:
                in_delim = 0
                continue
            elif char == "\001" and in_delim == 0:
                in_delim = 1
                continue
            elif char == ' ' and in_delim == 0:
                # print "%s %s" % (fields[fnum], f)
                da[fields[fnum]] = f
                fnum = fnum + 1
                if fnum == len(fields): break
                f = ""
                continue
            else: 
                f += char
                
        # Noticed that Kismet can send broken messages, so explicitly finishing fields to None here to debug:
        while fnum < len(fields):
            da[fields[fnum]] = None
            fnum = fnum + 1

        return da
    
    def getKismetProtocolTree(self, protocols): 
        for protocol in protocols:
            self.sendraw("!%i CAPABILITY %s\n" % (self.cmd, protocol.upper()))
            self.cmd = self.cmd + 1
            self.getKismetMessage()

    def subscribe(self, proto, fields, hndl, ackback = None):
        p = proto.upper()

        # print "subscribe %s %s" (p, fields)
        
        if p in ['ERROR', 'TIME', 'ACK', 'KISMET', 'PROTOCOLS', 'CAPABILITY']:
            raise KismetError(self.cmd, 'Cannot subscribe to protocol %s because it is handled by default' % p)

        if p in self.callbacks and self.callbacks[p] != None:
            print "!!! %s already declared with fields %s, replacing" % (p, self.callbacks[p][0])

        self.callbacks[p] = [ fields, hndl ]

        fe = ','.join(fields)

        if ackback != None: 
            self.ackbacks[self.cmd] = ackback

        self.sendraw("!%i ENABLE %s %s\n" % (self.cmd, p, fe))

        self.cmd = self.cmd + 1

    def unsubscribe(self, proto):
        p = proto.upper()
     
        self.callbacks.pop(p)
        
        self.sendraw("!0 REMOVE %s\n" % (p))
        
class KismetError(Exception):
    
    def __init__(self, id, value):
        self.id = id
        self.value = value
        
    def __str__(self):
        return "Error on command %i, %s" (self.id, self.value)

# Example of a generic callback:
def genericcb(proto, text):
    print "%s %s" % (proto, text)

# Example of a generic ackback:
def genericdiecb(text):
    print "ACK received"
    global k
    k.kill()
    sys.exit(0)
    
if __name__ == '__main__':
    k = Kismet()
    print k
    print k.protocols
    k.start()
    try:
        while k.isAlive(): time.sleep(1)
    except KeyboardInterrupt:
        print 'Killing thread...'
        k.kill()
